package com.wande.dataplatform.common.utils;

import com.wande.dataplatform.common.enums.DataPlatformErrorCode;
import com.wande.dataplatform.common.exception.DataPlatformException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;

/**
 * Doris 连接管理器
 * 
 * @author ruoyi
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class DorisConnectionManager {

    private final DataSource dorisDataSource;

    /**
     * 获取 Doris 连接
     * 
     * @return 数据库连接
     */
    public Connection getConnection() {
        try {
            Connection connection = dorisDataSource.getConnection();
            log.debug("成功获取 Doris 连接");
            return connection;
        } catch (SQLException e) {
            log.error("获取 Doris 连接失败", e);
            throw new DataPlatformException(DataPlatformErrorCode.DORIS_CONNECTION_FAILED, e);
        }
    }

    /**
     * 测试 Doris 连接
     * 
     * @return 连接是否正常
     */
    public boolean testConnection() {
        Connection conn = null;
        Statement stmt = null;
        ResultSet rs = null;
        
        try {
            conn = getConnection();
            stmt = conn.createStatement();
            stmt.setQueryTimeout(5); // 5秒超时
            rs = stmt.executeQuery("SELECT 1");
            
            boolean result = rs.next() && rs.getInt(1) == 1;
            log.info("Doris 连接测试{}", result ? "成功" : "失败");
            return result;
            
        } catch (Exception e) {
            log.error("Doris 连接测试失败", e);
            return false;
        } finally {
            closeResources(rs, stmt, conn);
        }
    }

    /**
     * 关闭数据库资源
     * 
     * @param rs 结果集
     * @param stmt 语句对象
     * @param conn 连接对象
     */
    public void closeResources(ResultSet rs, Statement stmt, Connection conn) {
        try {
            if (rs != null && !rs.isClosed()) {
                rs.close();
            }
        } catch (SQLException e) {
            log.warn("关闭 ResultSet 失败", e);
        }
        
        try {
            if (stmt != null && !stmt.isClosed()) {
                stmt.close();
            }
        } catch (SQLException e) {
            log.warn("关闭 Statement 失败", e);
        }
        
        try {
            if (conn != null && !conn.isClosed()) {
                conn.close();
            }
        } catch (SQLException e) {
            log.warn("关闭 Connection 失败", e);
        }
    }

    /**
     * 关闭连接
     * 
     * @param conn 连接对象
     */
    public void closeConnection(Connection conn) {
        closeResources(null, null, conn);
    }

}
